from datetime import datetime
import pandas as pd
import os
from collections import defaultdict
import numpy as np

# ========================================================================================
# ▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼ 用户自定义配置区域 - 按需修改以下参数 ▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼
# ========================================================================================

# 1. QA对输出文件目录和文件名
QA_OUTPUT_DIR = r""
QA_OUTPUT_FILE = os.path.join(QA_OUTPUT_DIR, "QA记录.xlsx")

# 2. 查询结果保存目录
QUERY_RESULT_DIR = r""

# 3. 设备名称列表
DEVICE_NAMES = [
    "启动回路母线",
    "SFC旁路隔离开关",
    "静止变频器（SFC）",
    "SFC 输出旁路刀",
    "输入开关（闸刀）",
    "SFC输入隔离开关",
    "输入变压器",
    "电流互感器",
    "电压互感器",
    "变频器系统",
    "电缆接线端子",
    "开关柜",
    "技术供水泵抽屉开关",
    "公用配电盘",
    "变压器",
    "厂高变",
    "高压开关柜",
    "机组自用配电盘",
    "耦合电容器",
    "220kVGIS正母线",
    "500kVGIS",
    "GIS",
    "GIS设备",
    "避雷器",
    "隔离刀闸",
    "刀闸",
    "断路器",
    "主变高压侧开关",
    "500kV电缆头",
    "高压电缆",
    "SFC电缆终端",
    "电抗器电缆终端",
    "电缆终端",
    "主变高压侧电缆终端",
    "电缆线",
    "电缆",
    "500kV出线线夹",
    "起励装置",
    "灭磁开关",
    "交流进线开关",
    "励磁灭磁开关",
    "励磁变压器",
    "灭磁电阻",
    "定子",
    "转子",
    "定子接地保护",
    "发电（电动）机",
    "球阀油泵",
    "机组泵",
    "组合轴承",
    "主变低压侧避雷器",
    "主变高压侧套管",
    "套管",
    "主变压器",
    "电气制动隔离开关（断路器）",
    "出口断路器",
    "接地刀闸",
    "被拖动刀闸",
    "被拖动隔离开关",
    "换相刀闸操作电机",
    "换相隔离开关",
    "换相闸刀",
    "换相刀闸",
    "启动刀闸",
    "启动母线闸刀",
    "发电机开关",
    "主变压器系统",
    "上库变"
]


# 生成文件配置
def generate_file_configs():
    configs = {}

    # 为电气绝缘所缺陷库生成配置
    configs["电气绝缘.xlsx"] = []
    for device in DEVICE_NAMES:
        configs["电气绝缘.xlsx"].append({
            "query_conditions": {"设备名称": device},
            "user_question": f"我想了解{device}曾经发生过什么缺陷？",
            "selected_columns": ["KKS编码", "设备名称", "缺陷名称（含缺陷特征）", "缺陷原因", "处理方案及措施",
                                 "缺陷来源", "来源文件", "缺陷等级", "缺陷时间"],
            "report_format": "{device_name}是抽水蓄能电站重要组成设备，在电气绝缘专业，曾经发生过危急缺陷{critical_count}项、严重缺陷{severe_count}项、一般缺陷{general_count}项。其中危急缺陷有\n{critical_defects}\n如果想了解以上缺陷的详细原因、处理措施，请回复具体缺陷名称+缺陷原因/处理措施。如果想了解其它等级缺陷，请回复严重/一般缺陷详细信息。",
            "defect_level_column": "缺陷等级",  # 指定分类统计字段
            "category_columns": ["缺陷等级"]  # 新增：需要分类统计的列
        })

    return configs


# ========================================================================================
# ▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲ 配置区域结束 ▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲▲
# ========================================================================================

def count_defects_by_level(data, defect_level_column):
    """统计不同缺陷等级的数量"""
    defect_counts = defaultdict(int)

    if defect_level_column not in data.columns:
        return defect_counts

    for level in data[defect_level_column]:
        if pd.isna(level):
            continue
        level_str = str(level).strip()
        if "危急" in level_str:
            defect_counts["critical"] += 1
        elif "严重" in level_str:
            defect_counts["severe"] += 1
        elif "一般" in level_str:
            defect_counts["general"] += 1

    return defect_counts


def format_date(date_str):
    """尝试格式化日期字符串"""
    try:
        # 尝试解析不同格式的日期
        formats = ["%Y.%m.%d", "%Y/%m/%d", "%Y-%m-%d", "%Y.%m", "%Y/%m", "%Y-%m", "%Y"]
        for fmt in formats:
            try:
                dt = datetime.strptime(str(date_str), fmt)
                return dt.strftime("%Y.%m.%d")
            except:
                continue
        # 如果所有格式都失败，返回原始字符串
        return str(date_str)
    except:
        return str(date_str)


def format_critical_defects(data, defect_level_column, defect_name_column, plant_name_column, date_column):
    """格式化危急缺陷示例（输出所有危急缺陷）"""
    critical_defects = []

    if defect_level_column not in data.columns:
        return "无危急缺陷记录"

    # 筛选危急缺陷
    critical_data = data[data[defect_level_column].astype(str).str.contains("危急", na=False)]

    if critical_data.empty:
        return "无危急缺陷记录"

    # 输出所有危急缺陷
    for i, (_, row) in enumerate(critical_data.iterrows()):
        defect_name = clean_text(row.get(defect_name_column, "未知缺陷"))
        plant_name = clean_text(row.get(plant_name_column, "未知单位"))
        date = clean_text(row.get(date_column, "未知日期"))

        # 格式化日期
        formatted_date = format_date(date)

        critical_defects.append(f"{i + 1}.{defect_name}（{plant_name}，{formatted_date}）")

    return "\n".join(critical_defects)


def clean_text(text) -> str:
    """清理文本数据"""
    if pd.isna(text):
        return "未知"
    if isinstance(text, str):
        return text.replace('\n', ' ').replace('\r', '').replace('\t', ' ').strip()
    return str(text)


def create_custom_report(row: dict, index: int, report_format: str, selected_columns: list) -> str:
    """根据配置的格式创建自定义报告"""
    # 准备数据
    format_data = {
        "序号": index,
        **{col: clean_text(row.get(col, "未知")) for col in selected_columns}
    }

    # 使用配置格式创建报告
    try:
        return report_format.format(**format_data)
    except KeyError as e:
        # 处理缺失的列
        missing = str(e).strip("'")
        format_data[missing] = clean_text(row.get(missing, "未知"))
        return report_format.format(**format_data)
    except:
        # 后备方案：如果格式错误，创建简单报告
        report_parts = [f"{index}."]
        for col in selected_columns:
            if col in row:
                report_parts.append(f"{col}:{clean_text(row[col])}")
        return " ".join(report_parts)


def query_equipment(data, conditions: dict, category_columns: list = None):
    """
    全字段模糊匹配设备缺陷查询函数
    新增功能：对指定列进行分类统计
    """
    result = data.copy()

    # 应用所有查询条件
    for field, pattern in conditions.items():
        if pattern and field in result.columns:
            try:
                # 对数值字段先转换为字符串
                if pd.api.types.is_numeric_dtype(result[field]):
                    result = result[result[field].astype(str).str.contains(str(pattern), na=False, case=False)]
                else:
                    result = result[result[field].astype(str).str.contains(str(pattern), na=False, case=False)]
            except Exception as e:
                print(f"警告: {field}字段处理错误: {str(e)}")
                continue

    # 新增：对指定列进行分类统计
    category_stats = {}
    if category_columns:
        for col in category_columns:
            if col in result.columns:
                # 统计每个类别的数量
                stats = result[col].value_counts().to_dict()

                # 将统计结果转换为字符串表示
                stats_str = ", ".join([f"{k}: {v}" for k, v in stats.items()])
                category_stats[col] = stats_str
                print(f"📊 {col}分类统计: {stats_str}")
            else:
                print(f"⚠️ 分类统计列不存在: {col}")

    return result, category_stats


def save_query_result(data, file_type: str):
    """保存查询结果到Excel"""
    if data.empty:
        print("⚠️ 无数据可保存")
        return None

    try:
        # 确保目录存在
        os.makedirs(QUERY_RESULT_DIR, exist_ok=True)

        # 自动生成文件名：知识库_时间戳.xlsx
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"{file_type}_查询结果_{timestamp}.xlsx"
        save_path = os.path.join(QUERY_RESULT_DIR, filename)

        # 保存数据
        data.to_excel(save_path, index=False)
        print(f"✅ 成功保存 {len(data)} 条记录到: {save_path}")
        return save_path

    except Exception as e:
        print(f"❌ 保存失败: {str(e)}")
        return None


def generate_qa_excel(query_result_file: str, config: dict):
    """生成QA对Excel文件并追加到固定文件"""
    try:
        # 1. 读取查询结果文件
        if not os.path.exists(query_result_file):
            print(f"❌ 错误：文件不存在 - {query_result_file}")
            return None

        source_df = pd.read_excel(query_result_file)
        print(f"✅ 成功读取查询结果数据，记录数: {len(source_df)}")

        # 2. 使用预设的列配置
        selected_columns = config["selected_columns"]
        report_format = config["report_format"]
        user_question = config["user_question"]
        defect_level_column = config.get("defect_level_column", "缺陷等级")

        valid_columns = [col for col in selected_columns if col in source_df.columns]
        missing_columns = [col for col in selected_columns if col not in source_df.columns]

        if not valid_columns:
            print("❌ 错误：选择的列在数据文件中都不存在")
            print(f"可用列: {', '.join(source_df.columns)}")
            return None

        if missing_columns:
            print(f"⚠️ 注意：以下列不存在于文件中: {', '.join(missing_columns)}")

        selected_df = source_df[valid_columns].copy()
        print(f"🟢 已选择列: {', '.join(valid_columns)}")

        # 3. 清理文本数据
        for col in valid_columns:
            selected_df[col] = selected_df[col].apply(clean_text)

        # 4. 统计缺陷等级数量
        defect_counts = count_defects_by_level(source_df, defect_level_column)

        # 5. 格式化危急缺陷示例
        critical_defects = format_critical_defects(
            source_df,
            defect_level_column,
            "缺陷名称（含缺陷特征）",
            "电站名称",
            "缺陷时间"
        )

        # 6. 创建自定义报告
        device_name = config["query_conditions"]["设备名称"]
        report = report_format.format(
            device_name=device_name,
            critical_count=defect_counts.get("critical", 0),
            severe_count=defect_counts.get("severe", 0),
            general_count=defect_counts.get("general", 0),
            critical_defects=critical_defects
        )

        # 7. 准备参考内容（原始数据）
        matching_rows = [row.to_dict() for _, row in selected_df.iterrows()]
        reference_lines = []
        for row in matching_rows:
            ref_parts = [f"{col}:{value}" for col, value in row.items()]
            reference_lines.append(", ".join(ref_parts))
        reference_content = "\n".join(reference_lines)

        # 8. 处理来源文件
        all_documents = set()
        if '来源文件' in valid_columns:
            all_documents = set(row.get('来源文件', '') for row in matching_rows)

        documents_str = ", ".join(filter(None, all_documents)) or "无记录"

        # 9. 创建最终记录
        os.makedirs(QA_OUTPUT_DIR, exist_ok=True)

        # 创建新记录
        record = {
            "id": len(matching_rows) + 1,
            "user_input": user_question,
            "reference": report,
            "reference_content": reference_content,
            "documents": documents_str,
            "response": "",  # 第六字段留空
            "retrieved_contexts": "[]"  # 第七字段为空列表
        }

        # 追加到输出文件
        if os.path.exists(QA_OUTPUT_FILE):
            existing_df = pd.read_excel(QA_OUTPUT_FILE)
            combined_df = pd.concat([existing_df, pd.DataFrame([record])], ignore_index=True)
        else:
            combined_df = pd.DataFrame([record])

        combined_df.to_excel(QA_OUTPUT_FILE, index=False)
        print(f"💾 成功追加QA对到文件: {QA_OUTPUT_FILE}")

        # 打印摘要
        print("=" * 80)
        print(f"📋 报告内容摘要:")
        print(f" - 问题: {user_question}")
        print(f" - 拼接报告: {report[:100]}{'...' if len(report) > 100 else ''}")
        print(f" - 包含记录: {len(matching_rows)}条")
        print(f" - 来源文件: {documents_str}")
        print("=" * 80)
        return QA_OUTPUT_FILE

    except Exception as e:
        print(f"❌ 生成QA对时出错: {str(e)}")
        return None


def process_file(file_path: str, config: dict):
    """处理单个Excel文件 - 使用预设配置"""
    file_name = os.path.basename(file_path)

    try:
        print(f"\n{'=' * 100}")
        print(f"🚀 开始处理文件: {file_name}")
        print("=" * 100)

        print(f"🔍 正在处理文件: {file_name}")

        # 1. 读取Excel文件并添加来源文件列
        try:
            # 尝试使用默认引擎读取Excel
            df = pd.read_excel(file_path)
        except Exception as e:
            # 如果是xlsb格式，尝试使用pyxlsb引擎
            if "Workbook is encrypted" in str(e) or "requires the pyxlsb package" in str(e):
                try:
                    print("⚠️ 检测到xlsb格式文件，尝试使用pyxlsb引擎...")
                    df = pd.read_excel(file_path, engine='pyxlsb')
                except ImportError:
                    print("❌ 需要安装pyxlsb库才能读取xlsb格式文件，请执行: pip install pyxlsb")
                    return False
                except Exception as inner_e:
                    print(f"❌ 使用pyxlsb引擎读取文件失败: {str(inner_e)}")
                    return False
            else:
                print(f"❌ 读取文件失败: {str(e)}")
                return False

        # 添加来源文件列，使用完整的文件名（带扩展名）
        df['来源文件'] = file_name
        print(f"✅ 已添加来源文件列: {file_name}")

        # 2. 输出配置信息
        print("\n配置文件:")
        print(f" - 查询条件: {config['query_conditions']}")
        print(f" - 用户问题: {config['user_question']}")
        print(f" - 选择列: {', '.join(config['selected_columns'])}")
        print(f" - 答案格式: {config['report_format']}")

        # 获取分类统计列
        category_columns = config.get("category_columns", [])
        if category_columns:
            print(f" - 分类统计列: {', '.join(category_columns)}")

        # 3. 执行查询
        query_results, category_stats = query_equipment(df, config["query_conditions"], category_columns)

        if query_results.empty:
            print("⚠️ 未找到匹配记录")
            return False

        # 4. 保存查询结果
        print(f"\n📊 找到 {len(query_results)} 条匹配记录")
        query_result_file = save_query_result(query_results, file_name.replace('.xlsx', ''))  # 保存时使用不带扩展名的文件名
        if not query_result_file:
            print("⚠️ 查询结果保存失败")
            return False

        # 5. 生成QA对
        qa_result = generate_qa_excel(query_result_file, config)

        if qa_result:
            print(f"✅ {file_name}处理完成")
            return True
        else:
            print(f"⚠️ {file_name}QA对生成失败")
            return False

    except FileNotFoundError:
        print(f"❌ 文件不存在: {file_path}")
        return False
    except Exception as e:
        print(f"❌ 处理文件时出错: {str(e)}")
        return False


def main():
    """主函数 - 依次处理所有预设文件"""
    # 定义文件路径
    source_dir = r""
    files = [
        os.path.join(source_dir, "四库建设-缺陷库.xlsx"),
        os.path.join(source_dir, "电气绝缘.xlsx"),
        os.path.join(source_dir, "金属缺陷库.xlsx")
    ]

    # 生成文件配置
    FILE_CONFIGS = generate_file_configs()

    success_count = 0
    total_configs = 0

    # 处理每个文件
    for file_path in files:
        file_name = os.path.basename(file_path)

        # 检查文件是否存在
        if not os.path.exists(file_path):
            print(f"❌ 文件不存在: {file_path}")
            continue

        # 获取该文件的配置列表
        if file_name in FILE_CONFIGS:
            configs = FILE_CONFIGS[file_name]
            total_configs += len(configs)

            print(f"\n{'=' * 100}")
            print(f"📁 开始处理文件: {file_name} (包含 {len(configs)} 个设备配置)")
            print("=" * 100)

            # 处理每个配置项
            for config in configs:
                print(f"\n🔧 处理设备: {config['query_conditions']['设备名称']}")
                if process_file(file_path, config):
                    success_count += 1
        else:
            print(f"⚠️ 没有为 {file_name} 找到配置，跳过处理")

    print("\n" + "=" * 100)
    print(f"🏁 程序执行完毕: 成功处理 {success_count}/{total_configs} 个设备配置")
    print("=" * 100)


if __name__ == "__main__":
    main()